"""
通过gzip和numpy解析MNIST数据集的二进制文件
"""

import os
import gzip
from cProfile import label

import numpy as np
import matplotlib.pyplot as plt


# 以n，28，28的形式读入mnist数据集, 或者读入标签。取决于文件名
def parse_mnist(mnist_file_addr: str = None, flatten: bool = False, one_hot: bool = False) -> np.array:
    """解析MNIST二进制文件, 并返回解析结果
    输入参数:
        minst_file: MNIST数据集的文件地址. 类型: 字符串.
        flatten: bool, 默认Fasle. 是否将图片展开, 即(n张, 28, 28)变成(n张, 784)
        one_hot: bool, 默认Fasle. 标签是否采用one hot形式.

    返回值:
        解析后的numpy数组
    """
    if mnist_file_addr is not None:
        mnist_file_name = os.path.basename(mnist_file_addr)  # 根据地址获取MNIST文件名字
        with gzip.open(filename=mnist_file_addr, mode="rb") as minst_file:
            mnist_file_content = minst_file.read()
        if "label" in mnist_file_name:  # 传入的为标签二进制编码文件地址
            data = np.frombuffer(buffer=mnist_file_content, dtype=np.uint8, offset=8)  # MNIST标签文件的前8个字节为描述性内容，直接从第九个字节开始读取标签，并解析
            if one_hot:
                data_zeros = np.zeros(shape=(data.size, 10))
                for idx, label in enumerate(data):
                    data_zeros[idx, label] = 1
                data = data_zeros
        else:  # 传入的为图片二进制编码文件地址
            data = np.frombuffer(buffer=mnist_file_content, dtype=np.uint8, offset=16)  # MNIST图片文件的前16个字节为描述性内容，直接从第九个字节开始读取标签，并解析
            data = data.reshape(-1, 784) if flatten else data.reshape(-1, 28, 28)

    return data


def distance(a, b):
    return np.sqrt(np.sum(np.square(a - b)))


class KNN:
    def __init__(self, k, label_num):
        self.k = k
        self.label_num = label_num

    def fit(self, data_train, label_train):
        self.data_train = data_train
        self.label_train = label_train

    def get_knn_indices(self, x):
        # 获取离某个点最近的k个点的索引值
        dis = []
        for i in range(0, len(self.data_train)):
            dis.append(distance(self.data_train[i], x))
        knn_indices = np.argsort(dis)
        return knn_indices[:self.k]

    def get_label(self, x):
        # 获取某个点x对应的标签
        knn_indices = self.get_knn_indices(x)
        label_statistic = np.zeros(shape=[self.label_num])
        for index in knn_indices:
            label = int(self.label_train[index])
            label_statistic[label] += 1
        return np.argmax(label_statistic)

    def predict(self, x_test):
        # 预测样本 test_x的类别
        predicted_test_labels = np.zeros(shape=[len(x_test)], dtype=int)
        for i in range(0, len(x_test)):
            predicted_test_labels[i] = self.get_label(x_test[i])
        return predicted_test_labels

if __name__ == "__main__":
    data_train = parse_mnist(mnist_file_addr="C:\\Users\ASUS\Desktop\机器学习\数据集\MNIST\\train-images-idx3-ubyte.gz")
    label_train = parse_mnist(mnist_file_addr="C:\\Users\ASUS\Desktop\机器学习\数据集\MNIST\\train-labels-idx1-ubyte.gz")

    data_test = parse_mnist(mnist_file_addr="C:\\Users\ASUS\Desktop\机器学习\数据集\MNIST\\t10k-images-idx3-ubyte.gz")
    label_test = parse_mnist(mnist_file_addr="C:\\Users\ASUS\Desktop\机器学习\数据集\MNIST\\t10k-labels-idx1-ubyte.gz")

    # data_train = data_train[:len(data_train) // 100]
    # label_train = label_train[:len(label_train) // 100]
    data_test = data_test[:len(data_test) // 500]
    label_test = label_test[:len(label_test) // 500]

    for k in range(1, 10):
        knn = KNN(k, label_num=10)
        knn.fit(data_train, label_train)
        predicted_labels = knn.predict(data_test)

        accuracy = np.mean(predicted_labels == label_test)
        print('when k is ' + str(k) + ', accuracy is ' + str(accuracy * 100)  + '%')
